package com.xyh.ollama_test.service;

import com.xyh.ollama_test.enums.Department;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.JsonWithInt;
import io.qdrant.client.grpc.Points;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;

@Service
public class BgeM3Service {
    private static final Logger logger = LoggerFactory.getLogger(BgeM3Service.class);
    private final WebClient webClient;
    private final QdrantClient qdrantClient;
    private static final String COLLECTION_NAME = "text_embeddings";
    private final PermissionService permissionService;

    public BgeM3Service(WebClient ollamaWebClient, QdrantClient qdrantClient, PermissionService permissionService) {
        this.webClient = ollamaWebClient;
        this.qdrantClient = qdrantClient;
        this.permissionService = permissionService;
    }

    public Mono<String> generateEmbeddingAndStore(String text, String userId) {
        return webClient.post()
                .uri("/api/embeddings")
                .bodyValue(Map.of(
                        "model", "bge-m3",
                        "prompt", text
                ))
                .retrieve()
                .bodyToMono(Map.class)
                .flatMap(response -> {
                    try {
                        logger.debug("Received embedding response: {}", response);

                        // 获取用户部门
                        Department userDept = permissionService.getUserDepartment(userId);
                        if (userDept == null) {
                            return Mono.error(new RuntimeException("用户部门信息不存在"));
                        }

                        List<Double> embeddingList = (List<Double>) response.get("embedding");
                        float[] floatArray = new float[embeddingList.size()];
                        for (int i = 0; i < embeddingList.size(); i++) {
                            floatArray[i] = embeddingList.get(i).floatValue();
                        }

                        Points.Vector vector = Points.Vector.newBuilder()
                                .addAllData(floatArrayToFloatList(floatArray))
                                .build();

                        Points.Vectors vectors = Points.Vectors.newBuilder()
                                .setVector(vector)
                                .build();

                        JsonWithInt.Value jsonValue = JsonWithInt.Value.newBuilder()
                                .setStringValue(text)
                                .build();

                        JsonWithInt.Value deptValue = JsonWithInt.Value.newBuilder()
                                .setStringValue(userDept.name())
                                .build();

                        Points.PointStruct point = Points.PointStruct.newBuilder()
                                .setId(Points.PointId.newBuilder()
                                        .setNum(UUID.randomUUID().getMostSignificantBits())
                                        .build())
                                .setVectors(vectors)
                                .putPayload("text", jsonValue)
                                .putPayload("department", deptValue)
                                .build();

                        CompletableFuture<Void> future = CompletableFuture.runAsync(() -> {
                            try {
                                qdrantClient.upsertAsync(COLLECTION_NAME, List.of(point)).get();
                                logger.info("Vector stored successfully for text: {}", text);
                            } catch (Exception e) {
                                logger.error("Failed to store vector for text: {}", text, e);
                                throw new RuntimeException("Failed to store vector", e);
                            }
                        });

                        return Mono.fromFuture(future)
                                .thenReturn("文本向量已成功存储: " + text);
                    } catch (Exception e) {
                        logger.error("Error processing embedding for text: {}", text, e);
                        return Mono.error(new RuntimeException("处理文本向量时出错", e));
                    }
                });
    }

    private List<Float> floatArrayToFloatList(float[] array) {
        Float[] boxedArray = new Float[array.length];
        for (int i = 0; i < array.length; i++) {
            boxedArray[i] = array[i];
        }
        return List.of(boxedArray);
    }
}